
#include <linux/module.h>

#include <linux/slab.h>
#include <linux/fs.h>
#include <linux/cdev.h>
#include <linux/device.h>


#define MEMDEV_NR_DEVS          3

static int mem_major;
module_param(mem_major, int, S_IRUGO);

#define MEMDEV_IOC_MAGIC        'M'
#define MEMDEV_IOC_CLEAR        _IO(MEMDEV_IOC_MAGIC, 1)
#define MEMDEV_IOC_FILL         _IOW(MEMDEV_IOC_MAGIC, 2, char)
#define MEMDEV_IOC_MAXNR        2

struct mem_device {
    struct cdev chdev;
    struct device dev;

    void *data;
    size_t size;
};

static int mem_open(struct inode *inode, struct file *filp)
{
    struct cdev *chdev = inode->i_cdev;
    struct mem_device *mdev = container_of(chdev, struct mem_device, chdev);

    filp->private_data = mdev;
    return 0;
}

static ssize_t mem_read(struct file *filp, char __user *buf, size_t count, loff_t *ppos)
{
    int ret;
    struct mem_device *dev = filp->private_data;

    if (*ppos >= dev->size) {
        return 0;
    }
    if (count > dev->size) {
        count = dev->size;
    }
    ret = copy_to_user(buf, dev->data + *ppos, count);
    if (ret) {
        ret = -EFAULT;
    } else {
        *ppos += count;
        ret = count;
    }

    return ret;
}

static ssize_t mem_write(struct file *filp, const char __user *buf, size_t count, loff_t *ppos)
{
    int ret;
    struct mem_device *dev = filp->private_data;

    if (*ppos >= dev->size) {
        return 0;
    }

    if (count + *ppos > dev->size) {
        count = dev->size - *ppos;
    }

    ret = copy_from_user(dev->data + *ppos, buf, count);
    if (ret) {
        ret = -EFAULT;
    } else {
        *ppos += count;
        memset(dev->data + *ppos, 0, dev->size - *ppos);
        ret = count;
    }

    return ret;
}

static long mem_ioctl(struct file *filp, unsigned int cmd, unsigned long arg)
{
    int ret = 0;
    struct mem_device *dev = filp->private_data;
    if (_IOC_TYPE(cmd) != MEMDEV_IOC_MAGIC) {
        return -ENOTTY;
    }
    if (_IOC_NR(cmd) > MEMDEV_IOC_MAXNR) {
        return -ENOTTY;
    }

    switch (cmd) {
    case MEMDEV_IOC_CLEAR:
        memset(dev->data, 0x00, dev->size);
        filp->f_pos = 0;
        break;

    case MEMDEV_IOC_FILL:
        memset(dev->data, (char)arg, dev->size);
        filp->f_pos = dev->size;
        break;

    default:
        ret = -EINVAL;
        break;
    }

    return ret;
}

static const struct file_operations mem_fops = {
    .owner      = THIS_MODULE,
    .open       = mem_open,
    .read       = mem_read,
    .write      = mem_write,
    .llseek     = default_llseek,
    .unlocked_ioctl = mem_ioctl,
    .compat_ioctl   = compat_ptr_ioctl,
};

struct class *memcls;
struct mem_device *memdevs[MEMDEV_NR_DEVS];

static void dev_release(struct device *dev)
{

}

static struct mem_device *mdev_init(struct class *cls, dev_t devno, int i)
{
    int ret;
    struct mem_device *mdev = kzalloc(sizeof(*mdev), GFP_KERNEL);
    if (mdev) {
        cdev_init(&mdev->chdev, &mem_fops);
        mdev->chdev.owner = THIS_MODULE;
        ret = cdev_add(&mdev->chdev, devno + i, 1);
        if (ret) {
            pr_err("cdev_add failed\n");
            goto fail;
        }

        mdev->dev.class = cls;
        mdev->dev.devt = devno + i;
        mdev->dev.release = dev_release;
        dev_set_name(&mdev->dev, "%s%d", KBUILD_MODNAME, i);

        ret = device_register(&mdev->dev);
        if (ret < 0) {
            pr_err("device_register failed\n");
            goto fail1;
        }

        mdev->size = 4 * 1024;
        mdev->data = (void *)__get_free_pages(GFP_KERNEL | __GFP_ZERO, get_order(mdev->size));
        // mdev->data = kzalloc(mdev->size, GFP_KERNEL);
        pr_debug("mdev->data: %px\n", mdev->data);
        if (IS_ERR_OR_NULL(mdev->data)) {
            pr_err("memory alloc failed\n");
            goto fail2;
        }
    }

    return mdev;
fail2:
    device_destroy(mdev->dev.class, mdev->dev.devt);
fail1:
    cdev_del(&mdev->chdev);
fail:
    kfree(mdev);
    return NULL;
}

static void mdev_deinit(struct mem_device *mdev)
{
    if (mdev->data) {
        free_pages((unsigned long)mdev->data, get_order(mdev->size));
        // kfree(mdev->data);
    }
    device_destroy(mdev->dev.class, mdev->dev.devt);
    cdev_del(&mdev->chdev);
    kfree(mdev);
}

static int __init memdev_init(void)
{
    int ret;
    int i;
    dev_t devno;

    if (mem_major) {
        devno = MKDEV(mem_major, 0);
        ret = register_chrdev_region(devno, MEMDEV_NR_DEVS, "mem-device");
    } else {
        ret = alloc_chrdev_region(&devno, 0, MEMDEV_NR_DEVS, "mem-device");
        mem_major = MAJOR(devno);
    }

    if (ret < 0) {
        return ret;
    }

    memcls = class_create(THIS_MODULE, "mem-device");

    for (i = 0; i < MEMDEV_NR_DEVS; i++) {
        struct mem_device *mdev = mdev_init(memcls, devno, i);
        if (!mdev) {
            int j;
            for (j = 0; j < i; j++) {
                mdev_deinit(memdevs[j]);
                memdevs[j] = NULL;
            }
            ret = -ENOMEM;
            goto fail;
        }
        memdevs[i] = mdev;
    }
    return 0;

fail:
    class_destroy(memcls);
    unregister_chrdev_region(devno, MEMDEV_NR_DEVS);
    return ret;
}


static void __exit memdev_exit(void)
{
    int i;
    for (i = 0; i < MEMDEV_NR_DEVS; i++) {
        struct mem_device *mdev = memdevs[i];
        mdev_deinit(mdev);
        memdevs[i] = NULL;
    }
    class_destroy(memcls);
    unregister_chrdev_region(MKDEV(mem_major, 0), MEMDEV_NR_DEVS);
}

module_init(memdev_init);
module_exit(memdev_exit);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("kingdix10");
